
package org.jeecg.modules.airag.llm.tools;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.jeecg.modules.airag.flow.component.enhance.IAiRagEnhanceJava;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.sql.*;
import java.util.*;

@Slf4j
@Component("generateDDL")
public class generateDDL implements IAiRagEnhanceJava {
    private final DataSource dataSource;
    private Map<String, String> tableStructures; // 这个字段未初始化

    public generateDDL(DataSource dataSource) {
        this.dataSource = dataSource;
        this.tableStructures = new HashMap<>(); // 初始化tableStructures
    }

    @Override
    public Map<String, Object> process(Map<String, Object> map) {
        List<String> tableNames = new ArrayList<>();
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            Object resultObj = map.get("result");

            JsonNode jsonNode;
            if (resultObj instanceof String) {
                // 如果是字符串，直接解析
                jsonNode = objectMapper.readTree((String) resultObj);
            } else if (resultObj instanceof List) {
                // 如果是List，将其转换为JsonNode
                jsonNode = objectMapper.valueToTree(resultObj);
            } else {
                throw new IllegalArgumentException("result must be either String or List, but got: " +
                        (resultObj != null ? resultObj.getClass() : "null"));
            }
            StringBuilder sb = new StringBuilder();
            for (JsonNode node : jsonNode) {
                // 提取 main_table 和 detail_table 的值
                String mainTable = node.get("main_table").asText();
                String detailTable = node.get("detail_table").asText();

                // 将值添加到 StringBuilder 中，用逗号分隔
                sb.append(mainTable).append(",").append(detailTable).append(",");
            }
            // 移除末尾的逗号并分割成列表
            if (sb.length() > 0) {
                String tableStr = sb.substring(0, sb.length() - 1); // 移除最后一个逗号
                String[] tables = tableStr.split(",");
                tableNames = Arrays.asList(tables);
            }

        } catch (Exception e) {
            log.error("解析result时发生错误", e);
            throw new RuntimeException("解析result时发生错误", e);
        }

        String ddl = generateRelevantDDL(tableNames);
        log.info("成功生成表结构DDL：{}", ddl);
        return Map.of("result", ddl);
    }




    /**
     * 生成数据库的DDL语句
     * 遍历数据库中所有以'py'开头的表，生成它们的建表语句
     *
     * @param dataSource 数据源对象
     * @return 包含所有表DDL语句的字符串
     * @throws RuntimeException 当生成DDL过程中发生SQL异常时
     */
    private static String generateDDL(DataSource dataSource) {
        log.info("开始生成数据库DDL");
        StringBuilder ddl = new StringBuilder();
        try (Connection connection = dataSource.getConnection()) {
            DatabaseMetaData metaData = connection.getMetaData();
            String catalog = connection.getCatalog();
            ResultSet tables = metaData.getTables(catalog, null, "py%", new String[]{"TABLE"});
            log.info("成功获取数据库连接和元数据信息");

            while (tables.next()) {
                String tableName = tables.getString("TABLE_NAME");
                // 跳过系统表
                if (tableName.startsWith("sys_") || tableName.startsWith("information_schema")) {
                    log.info("跳过系统表: {}", tableName);
                    continue;
                }
                log.info("正在处理表: {}", tableName);
                String createTableStatement = generateCreateTableStatement(tableName, metaData);
                ddl.append(createTableStatement).append("\n");
            }
        } catch (SQLException e) {
            log.error("生成DDL时发生错误", e);
            throw new RuntimeException(e);
        }
        log.info("DDL生成完成");
        return ddl.toString();
    }


    private String generateRelevantDDL(List<String> tableNames) {
        StringBuilder ddl = new StringBuilder();
        Map<String, String> localTableStructures = tableStructures != null ? tableStructures : new HashMap<>();

        for (String tableName : tableNames) {
            tableName = tableName.trim();
            if (localTableStructures.containsKey(tableName)) {
                ddl.append(localTableStructures.get(tableName)).append("\n");
                continue;
            }

            try (Connection connection = dataSource.getConnection()) {
                DatabaseMetaData metaData = connection.getMetaData();
                String createTableStatement = generateCreateTableStatement(tableName, metaData);
                localTableStructures.put(tableName, createTableStatement);
                ddl.append(createTableStatement).append("\n");
            } catch (SQLException e) {
                log.error("生成表 {} 的DDL时发生错误", tableName, e);
            }
        }

        // 更新实例变量
        this.tableStructures = localTableStructures;
        return ddl.toString();
    }

    /**
     * 生成单个表的建表语句
     * 包含列名、类型、是否可空、主键和注释等信息
     *
     * @param tableName 表名
     * @param metaData 数据库元数据对象
     * @return 表的建表语句
     * @throws SQLException 当获取表结构信息过程中发生SQL异常时
     */
    private static String generateCreateTableStatement(String tableName, DatabaseMetaData metaData) throws SQLException {
        log.info("开始生成表 {} 的建表语句", tableName);
        try (Connection connection = metaData.getConnection();
             Statement statement = connection.createStatement();
             ResultSet resultSet = statement.executeQuery("SHOW CREATE TABLE " + tableName)) {

            if (resultSet.next()) {
                // 获取CREATE TABLE语句（在MySQL中是第二列）
                String createTableStatement = resultSet.getString(2);
                log.debug("成功获取表 {} 的建表语句", tableName);
                log.info("成功生成表 {} 的建表语句", createTableStatement);
                return createTableStatement + "\n";
            } else {
                log.warn("无法获取表 {} 的建表语句", tableName);
                return "";
            }
        } catch (SQLException e) {
            log.error("获取表 {} 的建表语句时发生错误", tableName, e);
            throw e;
        }
    }
}
